import tensorflow as tf
from tensorflow.keras.layers import *
from PIL import Image
import numpy as np

class BasicImageClassification(object):
    def __init__(self,kinds) -> None:
        super().__init__()
        self.TFSeq=tf.keras.Sequential()
        if type(kinds) != list:
            raise Exception('Kinds MUST be a list with names')
        self.TrainDatas=[]
        self.Kinds=kinds
        self.TrainKinds=[]
        self.KindsLen=len(self.Kinds)
    def FeedImage(self,imagePath,kind):
        with Image.open(imagePath) as im:
            im=im.convert('L')
            im=im.resize((128,128))
            self.TrainDatas.append(np.array(im))
            print(im.size)
            k=np.asarray([0]*self.KindsLen)
            k[kind]=1
            self.TrainKinds.append(k)
    def BuildUp(self):
        #输入层
        #预处理
        self.TFSeq.add(Flatten(input_shape=(128,128)))
        #隐藏层
        self.TFSeq.add(Dense(256,activation='tanh'))
        self.TFSeq.add(Dropout(0.2))
        self.TFSeq.add(Dense(64,activation='relu'))
        #输出层
        self.TFSeq.add(Dense(self.KindsLen))
        self.TFSeq.compile(optimizer='adam',loss='mse',metrics=['accuracy'])
    def Train(self):
        self.TFSeq.summary()
        self.TFSeq.fit(np.asarray(self.TrainDatas),np.asarray(self.TrainKinds),epochs=10)
